from refer_path import ReferPath
from spline import ConstVelocitySpline
from mpc_cvxpy import MPC
from model import Model
import matplotlib.pyplot as plt

from celluloid import Camera

line = ConstVelocitySpline(
    [(0, 0), (1, 0), (2, 1), (3, 0), (4, -1), (5, 0), (-5, 0)], 1, 0.02
)
path = ReferPath(line)
model = Model((-100, 100), (-100, 100), (-100, 100))
mpc = MPC(model, path, T=10, timeStep=0.1, Nx=4, Nu=2, Q=[10, 1, 100, 100], R=[1, 1])

fig = plt.figure()
camera = Camera(fig)

# plt.scatter(line.x, line.y,color='blue')

points = []
for i in range(200):
    points.append(model.getState())
    ctrl, states, controls = mpc.getControl(model.getState())
    model.update(ctrl[0], ctrl[1], 0.1)

    plt.plot(line.x, line.y, color="blue")
    mpc.drawPrediction(plt, states, "red")
    plt.plot([x[2] for x in points], [x[3] for x in points], color="green")
    plt.scatter([x[2] for x in points], [x[3] for x in points], color="green")
    plt.pause(0.001)
    camera.snap()

animation = camera.animate()
animation.save("trajectory.gif")

plt.show()
